from os import environ
from typing import Type, Optional

from baidubce.auth.bce_credentials import BceCredentials
from baidubce.bce_client_configuration import BceClientConfiguration
from baidubce.exception import BceServerError, BceHttpClientError
from baidubce.services.bts.bts_client import BtsClient
from baidubce.services.bts.model import Row, Cell, QueryRowArgs, ScanArgs, BatchPutRowArgs
from flask import current_app

from app.models.TSModel import TSModel


class TableStoreClient:

    def __init__(self):
        self.bts = None

    def init_bts_client(self) -> None:
        config = BceClientConfiguration(
            credentials=BceCredentials(environ.get('BTS_ACCESS_KEY'), environ.get('BTS_SECRET_KEY')),
            endpoint=environ.get('BTS_ENDPOINT'),
        )
        self.bts = BtsClient(config)

    def db_is_not_none(self):
        if self.bts is None:
            raise RuntimeError("self.bts database instance  is None.")

    def batch_save(self, models: [TSModel]) -> bool:
        if models is None:
            raise RuntimeError("The provided models is None and cannot be saved to the database.")
        model_cls = None
        batch_put_row = BatchPutRowArgs()
        for model in models:
            fields = model.model_dump()
            row = Row()
            row.rowkey = model.get_row_key()
            row.cells.extend([Cell(name, value).__dict__ for name, value in fields.items()])
            batch_put_row.rows.append(row)
            model_cls = model.__class__
        return self.bts.batch_put_row(model_cls.__INSTANCE_NAME__, model_cls.__TABLE_NAME__, batch_put_row) is not None

    def save(self, model: TSModel) -> bool:
        if model is None:
            raise RuntimeError("The provided model instance is None and cannot be saved to the database.")
        fields = model.model_dump()
        row = Row()
        row.rowkey = model.get_row_key()
        row.cells.extend([Cell(name, value).__dict__ for name, value in fields.items()])
        return self.bts.put_row(model.__INSTANCE_NAME__, model.__TABLE_NAME__, row) is not None

    def query(self, row_key: str, model_cls: Type[TSModel]):
        query_row_args = QueryRowArgs()
        query_row_args.rowkey = row_key
        try:
            query_response = self.bts.get_row(model_cls.__INSTANCE_NAME__, model_cls.__TABLE_NAME__, query_row_args)
            data = next(iter(query_response.result), None) if query_response.result is not None else None
            return model_cls(**{cell.column: cell.value for cell in data.cells}) if data is not None else None
        except (BceServerError, BceHttpClientError) as ex:
            current_app.logger.error(f"{str(model_cls.__TABLE_NAME__)},rowkey:{row_key}",ex)
            return None

    def list(self, model_cls: Type[TSModel], start_row_key: Optional[str] = "", stop_rowkey: Optional[str] = ""):
        scan_args = ScanArgs(start_rowkey=start_row_key, stop_rowkey=stop_rowkey)
        response = self.bts.scan(model_cls.__INSTANCE_NAME__, model_cls.__TABLE_NAME__, scan_args)
        if response.result is not None:
            result_list = []
            for data in response.result:
                data_dict = {cell.column: cell.value for cell in data.cells}
                result_list.append(model_cls(**data_dict))
            return result_list
        else:
            return None


ts_client = TableStoreClient()
